package com.icinfo.cloud.provider.common.utils;

import com.icinfo.cloud.provider.common.exception.BusinessException;
import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.security.*;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

public class RSAEncryptUtil {

    private static final String RSA = "RSA";
    public static final String RSA_SHA256 = "SHA256withRSA";

    /**
     * RSA sign
     * alg:SHA256withRSA
     *
     * @param privateKey
     * @param data
     * @return
     * @throws Exception
     */
    public static String sign(final PrivateKey privateKey, final byte[] data) throws BusinessException {
        try {
            Signature signature = Signature.getInstance(RSA_SHA256);
            signature.initSign(privateKey);
            signature.update(data);
            return Base64.encodeBase64String(signature.sign().clone());

        } catch (Exception e) {
            e.printStackTrace();
            throw new BusinessException("RSA sign error");
        }
    }

    /**
     * RSA sign
     * alg:SHA256withRSA
     *
     * @param privateKey
     * @param data
     * @return
     * @throws Exception
     */
    public static String sign(final String privateKey, final String data) throws BusinessException {
        return sign(getPrivateKey(privateKey), data.getBytes());
    }

    /**
     * RSA verify
     *
     * @param publicKey
     * @param data
     * @param sign
     * @return
     * @throws Exception
     */
    public static boolean verify(final PublicKey publicKey, final byte[] data, final byte[] sign) throws Exception {
        Signature signature = Signature.getInstance(RSA_SHA256);
        signature.initVerify(publicKey);
        signature.update(data);
        return signature.verify(sign);
    }

    public static boolean verify(final String publicKey, final String data, final String sign) throws Exception {
        return verify(getPublicKey(publicKey), data.getBytes(), Base64.decodeBase64(sign));
    }

    /**
     * RSA encrypt
     *
     * @param publicKey
     * @param data
     * @return
     * @throws Exception
     */
    public static byte[] encrypt(final Key publicKey, final byte[] data) throws Exception {
        Cipher cipher = Cipher.getInstance(RSA);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] result = cipher.doFinal(data);
        return result;
    }

    /**
     * encrypt String
     *
     * @param publicKey
     * @param data
     * @return
     * @throws Exception
     */
    public static String encrypt(final String publicKey, final String data) throws Exception {
        Cipher cipher = Cipher.getInstance(RSA);
        cipher.init(Cipher.ENCRYPT_MODE, getPublicKey(publicKey));
        byte[] bytes = data.getBytes();
        int inputLen = bytes.length;
        int offLen = 0;
        int i = 0;
        ByteArrayOutputStream bops = new ByteArrayOutputStream();
        while (inputLen - offLen > 0) {
            byte[] cache;
            if (inputLen - offLen > 255) {
                cache = cipher.doFinal(bytes, offLen, 255);
            } else {
                cache = cipher.doFinal(bytes, offLen, inputLen - offLen);
            }
            bops.write(cache);
            i++;
            offLen = 255 * i;
        }
        bops.close();
        return Base64.encodeBase64String(bops.toByteArray());
    }

    /**
     * RSA decrypt
     *
     * @param privateKey
     * @param data
     * @return
     * @throws Exception
     */
    public static byte[] decrypt(final Key privateKey, final byte[] data) throws Exception {
        Cipher deCipher = Cipher.getInstance(RSA);
        deCipher.init(Cipher.DECRYPT_MODE, privateKey);
        byte[] result = deCipher.doFinal(data);
        return result;
    }

    /**
     * decrypt String
     *
     * @param privateKey
     * @param data
     * @return
     * @throws Exception
     */
    public static String decrypt(final String privateKey, final String data) throws Exception {
        Cipher deCipher = Cipher.getInstance(RSA);
        deCipher.init(Cipher.DECRYPT_MODE, getPrivateKey(privateKey));
        byte[] bytes = Base64.decodeBase64(data);
        int inputLen = bytes.length;
        int offLen = 0;
        int i = 0;
        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
        while (inputLen - offLen > 0) {
            byte[] cache;
            if (inputLen - offLen > 256) {
                cache = deCipher.doFinal(bytes, offLen, 256);
            } else {
                cache = deCipher.doFinal(bytes, offLen, inputLen - offLen);
            }
            byteArrayOutputStream.write(cache);
            i++;
            offLen = 256 * i;

        }
        byteArrayOutputStream.close();
        byte[] byteArray = byteArrayOutputStream.toByteArray();
        return new String(byteArray);
    }

    /**
     * get pubkey from cert
     *
     * @param file
     * @return
     * @throws Exception
     */
    public static PublicKey getPubKey(final String file) throws Exception {
        CertificateFactory cf = CertificateFactory.getInstance("X.509");
        X509Certificate cert = (X509Certificate) cf.generateCertificate(new FileInputStream(file));
        PublicKey publicKey = cert.getPublicKey();
        return publicKey;
    }

    /**
     * get public/private key pair from jks
     *
     * @param file
     * @param jksPassword
     * @param certPassword
     * @param certAlias
     * @return
     * @throws Exception
     */
    public static Map<String, Object> getJKSKeyPair(final String file, final String jksPassword,
                                                    final String certPassword, final String certAlias) throws Exception {
        KeyStore keyStore = KeyStore.getInstance("JKS");
        keyStore.load(new FileInputStream(file), jksPassword.toCharArray());
        PrivateKey privateKey = (PrivateKey) keyStore.getKey(certAlias, certPassword.toCharArray());
        PublicKey publicKey = keyStore.getCertificate(certAlias).getPublicKey();
        Map<String, Object> result = new HashMap<>();
        result.put("publicKey", publicKey);
        result.put("privateKey", privateKey);
        return result;
    }

    /**
     * get public key by String
     *
     * @param key String to public key
     * @return
     * @throws Exception
     */
    public static PublicKey getPublicKey(String key) throws BusinessException {
        try {
            byte[] keyBytes = Base64.decodeBase64(key);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA);
            PublicKey publicKey = keyFactory.generatePublic(keySpec);
            return publicKey;
        } catch (Exception e) {
            e.printStackTrace();
            throw new BusinessException("RSA getPublicKey error");
        }

    }

    /**
     * get private key by String
     *
     * @param key String to PrivateKey
     * @throws Exception
     */
    public static PrivateKey getPrivateKey(String key) throws BusinessException {
        try {
            byte[] keyBytes = Base64.decodeBase64(key);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA);
            PrivateKey privateKey = keyFactory.generatePrivate(keySpec);
            return privateKey;

        } catch (Exception e) {
            e.printStackTrace();
            throw new BusinessException("RSA getPrivateKey error");
        }
    }

}
